#!/usr/bin/python

import unittest

from hips.scheduling import constraint
from hips.scheduling import constraint_mocks
from hips.scheduling import schedule_mocks
from hips.scheduling import solution_log


def CallWithMockAsst(c, asst):
  return c(*constraint_mocks.CreateMockConstraintParameters(asst))


NONE = 0
ALICE = 1 << 0
BOB = 1 << 1
CARL = 1 << 2
DEE = 1 << 3
EARL = 1 << 4
FAYE = 1 << 5

DOCTOR_COUNT = 6


class ConstraintTestCase(unittest.TestCase):

  # NoTwoConsecutiveShifts tests

  def testNoTwoConsecutiveShiftsOk1(self):
    """NoTwoConsecutiveShifts, no conflicts, simple example."""
    a = [EARL | CARL, FAYE,
         ALICE | CARL, BOB,
         ALICE | CARL, EARL,
         ALICE | CARL, FAYE,
         BOB | CARL, DEE,
         BOB | FAYE, CARL,
         BOB | FAYE, EARL,
         DEE | BOB, ALICE,
         DEE | EARL, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.NoTwoConsecutiveShifts(
        DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(0, conflicts)

  def testNoTwoConsecutiveShiftsConflict1(self):
    """NoTwoConsecutiveShifts, 1 conflict, simple example."""
    a = [EARL | CARL, FAYE,
         ALICE | CARL, BOB,
         ALICE | CARL, EARL,
         ALICE | CARL, FAYE,
         BOB | CARL, DEE,
         BOB | FAYE, CARL,
         DEE | FAYE, BOB,
         DEE | BOB, ALICE,  # conflict: BOB just worked the previous night
         DEE | EARL, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.NoTwoConsecutiveShifts(
        DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(1, conflicts)

  # NoMoreThanMaxConsecutiveDays tests

  def testNoMoreThanMaxConsecutiveDaysOk1(self):
    """NoMoreThanMaxConsecutiveDays, no conflicts, basic example."""
    a = [EARL | CARL, FAYE,
         ALICE | CARL, BOB,
         ALICE | CARL, EARL,
         ALICE | CARL, FAYE,
         BOB | CARL, DEE,
         BOB | FAYE, CARL,
         BOB | FAYE, EARL,
         DEE | BOB, ALICE,
         DEE | EARL, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.NoMoreThanMaxConsecutiveDays(
        [4, 4, 6, 3, 3, 3],
        DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(0, conflicts)

  def testNoMoreThanMaxConsecutiveDaysConflicts1(self):
    """NoMoreThanMaxConsecutiveDays, 1 conflict."""
    a = [EARL | CARL, FAYE,
         ALICE | CARL, BOB,
         ALICE | CARL, EARL,
         ALICE | CARL, FAYE,
         BOB | CARL, DEE,
         BOB | FAYE, CARL,  # conflict: CARL's 6th day
         BOB | FAYE, EARL,
         DEE | BOB, ALICE,
         DEE | EARL, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.NoMoreThanMaxConsecutiveDays(
        [4, 4, 5, 3, 3, 3],
        DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(1, conflicts)

  # TotalShiftsInRange tests

  def testTotalShiftsInRangeOk1(self):
    """TotalShiftsInRange, no conflicts, basic example."""
    # This violates overlap constraint and min consecutive, but hey.
    a = [ALICE | BOB, FAYE,
         CARL | DEE, BOB,
         EARL | FAYE, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.TotalShiftsInRange(
        [1, 1, 1, 1, 1, 1],
        [2, 2, 1, 1, 1, 2],
        [1, 1, 1, 1, 1, 1],
        DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(0, conflicts)

  def testTotalShiftsInRangeOk2(self):
    """TotalShiftsInRange, no conflicts, bigger example."""
    # Wow it's annoying to come up with this test case!
    # I can see how real scheduling is a pain. :-)
    a = [EARL | CARL, FAYE,
         ALICE | CARL, BOB,
         ALICE | CARL, EARL,
         ALICE | CARL, FAYE,
         BOB | CARL, DEE,
         BOB | FAYE, CARL,
         BOB | FAYE, EARL,
         DEE | BOB, ALICE,
         DEE | EARL, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.TotalShiftsInRange(
        [4, 4, 4, 3, 4, 3],
        [6, 6, 6, 4, 5, 4],
        [2, 2, 2, 2, 2, 2],
        DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(0, conflicts)

  def testTotalShiftsInRangeConflict1(self):
    """TotalShiftsInRange, single conflicts, basic example."""
    # This violates overlap constraint and min consecutive, but hey.
    a = [ALICE | BOB, FAYE,
         CARL | DEE, FAYE,   # conflict: too many nights for FAYE
         EARL | FAYE, ALICE] # conflict: too many shifts for FAYE
    log = solution_log.SolutionLog()
    c = constraint.TotalShiftsInRange(
        [1, 1, 1, 1, 1, 1],
        [2, 2, 1, 1, 1, 2],
        [1, 1, 1, 1, 1, 1],
        DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(2, conflicts)

  # TODO Test NotInDay

  # NoLessThanMinConsecutiveDayShifts tests

  def testNoLessThanMinConsecutiveDayShiftsOk1(self):
    """NoLessThanMinConsecutiveDayShifts, no conflicts."""
    a = [ALICE | CARL, FAYE,
         ALICE | CARL, DEE,
         ALICE | CARL, EARL,
         ALICE | CARL, EARL,
         BOB | CARL, DEE,
         BOB | FAYE, EARL,
         BOB | FAYE, EARL,
         CARL | FAYE, BOB,
         CARL | FAYE, DEE,
         CARL | BOB, EARL,
         CARL | BOB, FAYE,
         DEE | BOB, ALICE,
         DEE | BOB, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.NoLessThanMinConsecutiveDayShifts(DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(0, conflicts)

  def testNoLessThanMinConsecutiveDayShiftsConflict1(self):
    """NoLessThanMinConsecutiveDayShifts, simple conflict."""
    a = [ALICE | CARL, FAYE,
         ALICE | CARL, DEE,
         ALICE | BOB, EARL,  # conflict
         ALICE | CARL, EARL,
         BOB | CARL, DEE,
         BOB | FAYE, EARL,
         BOB | FAYE, EARL,
         CARL | FAYE, BOB,
         CARL | FAYE, DEE,
         CARL | BOB, EARL,
         CARL | BOB, FAYE,
         DEE | BOB, ALICE,
         DEE | BOB, ALICE]
    log = solution_log.SolutionLog()
    c = constraint.NoLessThanMinConsecutiveDayShifts(DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    # Note this causes 3 conflicts: CC B CC each count as one.
    self.assertEquals(3, conflicts)

  # DayShiftsOverlap tests

  def testDayShiftsOverlapConstraintOk1(self):
    """DayShiftsOverlap, no conflicts."""
    # Note that night shifts don't matter here.
    a = [ALICE | BOB, FAYE,
         ALICE | CARL, DEE,
         CARL | FAYE, DEE,
         CARL | FAYE, ALICE,
         FAYE | EARL, BOB]
    log = solution_log.SolutionLog()
    c = constraint.DayShiftsOverlap(DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(0, conflicts)

  def testDayShiftsOverlapConstraintConflict1(self):
    """DayShiftsOverlap, single conflict."""
    # Note that night shifts don't matter here.
    a = [ALICE | BOB, FAYE,
         EARL | CARL, DEE,  # conflict
         CARL | FAYE, DEE,
         CARL | FAYE, ALICE,
         FAYE | EARL, BOB]
    log = solution_log.SolutionLog()
    c = constraint.DayShiftsOverlap(DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(1, conflicts)

  def testDayShiftsOverlapConstraintConflict2(self):
    """DayShiftsOverlap, multiple conflict."""
    # Note that night shifts don't matter here.
    a = [ALICE | BOB, FAYE,
         EARL | CARL, DEE,  # conflict
         ALICE | FAYE, DEE,  # conflict
         CARL | FAYE, ALICE,
         FAYE | EARL, BOB]
    log = solution_log.SolutionLog()
    c = constraint.DayShiftsOverlap(DOCTOR_COUNT,
        schedule_mocks.CreateMockSchedule(asst=a),
        log)
    conflicts = CallWithMockAsst(c, a)
    self.assertEquals(2, conflicts)


if __name__ == '__main__':
  unittest.main()
